{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specifying Custom Gradients\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\n",
    "\n",
    "Drake offers automatic differentiation (AutoDiff) capabilities to solve mathematical programs.\n",
    "However, there are some cases in which you will need to specify your own gradients to the optimization solver due to the problem requirements.\n",
    "\n",
    "This tutorial will cover two ways of providing gradients: *analytically*, and *numerically* using finite differencing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining a Simple Mathematical Program\n",
    "In this example, we are optimizing a simple 2-variable cost function:\n",
    "\n",
    "$\\begin{aligned}\n",
    "f(x) = x_0 + \\frac{3}{2} x_1\n",
    "\\end{aligned}$\n",
    "\n",
    "where both variables are constrained to be in the range `[0..1]`.\n",
    "\n",
    "Without any additional constraints, we can find the optimal solution to be `x_0 = x_1 = 0`, which yields a cost of `0`.\n",
    "\n",
    "By adding a constraint that the solution cannot lie on a set of circles:\n",
    " * One centered at the origin with radius `0.5`\n",
    " * One centered at `x=0.75, y=-0.15` with radius `0.25`\n",
    " * One centered at `x=0.55, y=0.06` with radius `0.1`\n",
    "\n",
    "The problem can formally be described as:\n",
    "\n",
    "$\\begin{aligned}\n",
    "\\min_x\\hspace{10pt} \\; & f(x) = x_0 + \\frac{3}{2} x_1 \\\\\n",
    "\\text{subject to}\\hspace{10pt} \\; & x_0^2 + x_1^2 >= 0.5^2 \\\\\n",
    " \\; & (x_0 - 0.75)^2 + (x_1 + 0.15)^2 >= 0.25^2 \\\\\n",
    " \\; & (x_0 - 0.55)^2 + (x_1 - 0.06)^2 >= 0.1^2\n",
    "\\end{aligned}$\n",
    "\n",
    "Here, we instead get an optimal solution of `x_0 = 0.48192` and `x_1 = 0.13325`, which yields a cost of `0.6818`.\n",
    "This is shown in the plot below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "x0 = np.linspace(0.0, 1.0, 11)\n",
    "x1 = np.linspace(0.0, 1.0, 11)\n",
    "X0, X1 = np.meshgrid(x0, x1)\n",
    "cost = X0 + 1.5 * X1\n",
    "cost_contour = ax.contourf(X0, X1, cost, levels=100)\n",
    "color_bar = plt.colorbar(cost_contour)\n",
    "color_bar.set_label(\"cost\")\n",
    "\n",
    "ax.add_patch(plt.Circle((0.0, 0.0), 0.5, color=\"white\"))\n",
    "ax.add_patch(plt.Circle((0.75, -0.15), 0.25, color=\"white\"))\n",
    "ax.add_patch(plt.Circle((0.55, 0.06), 0.1, color=\"white\"))\n",
    "\n",
    "ax.plot(0.48192, 0.13325, \"rx\", markersize=10, markeredgewidth=2)\n",
    "\n",
    "ax.set_xlim([0.0, 1.0])\n",
    "ax.set_ylim([0.0, 1.0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solving the Program\n",
    "As you have seen in the other tutorials, Drake allows you to create a `MathematicalProgram` to solve such a problem.\n",
    "\n",
    "For simple algebraic expressions like the ones in our example, Drake's automatic differentiation (or AutoDiff) will take care of computing all the necessary gradients for the optimization solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.solvers import MathematicalProgram, Solve\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "prog.AddCost(x[0] + 1.5 * x[1])\n",
    "prog.AddBoundingBoxConstraint(0.0, 1.0, x)\n",
    "prog.AddConstraint(x[0]**2 + x[1]**2 >= 0.5**2)\n",
    "prog.AddConstraint((x[0] - 0.75)**2 + (x[1] + 0.15)**2 >= 0.25**2)\n",
    "prog.AddConstraint((x[0] - 0.55)**2 + (x[1] - 0.06)**2 >= 0.1**2)\n",
    "\n",
    "result = Solve(prog, np.array([0.5, 0.5]))\n",
    "print(\"Success? \", result.is_success())\n",
    "print(result.get_solution_result())\n",
    "print(\"x* = \", result.GetSolution(x))\n",
    "print(\"cost = \", result.get_optimal_cost())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Providing Your Own Analytical Gradients\n",
    "For this specific example, we have already seen that providing gradients is not needed.\n",
    "\n",
    "However, if your computations are complex enough that they are intractable or slow via Drake's automatic differentiation, or they require a third-party library incompatible with automatic differentiation, you may need to provide the gradients yourself.\n",
    "\n",
    "In these cases, you can provide analytical gradients as shown below. For our example, we can do this because we can calculate the derivatives of the costs and constraints.\n",
    "\n",
    "$\\begin{aligned}\n",
    "\\frac{\\partial f(x)}{\\partial x} = [1, 1.5] \n",
    "\\end{aligned}$\n",
    "\n",
    "$\\begin{aligned}\n",
    "\\frac{\\partial \\text{constraint}}{\\partial x} = \\begin{bmatrix}\n",
    "2x_0 & 2x_1 \\\\\n",
    "2(x_0 - 0.75) & 2(x_1 + 0.15) \\\\\n",
    "2(x_0 - 0.55) & 2(x_1 - 0.06)\n",
    "\\end{bmatrix}\n",
    "\\end{aligned}$\n",
    "\n",
    "**NOTE:** It is worth reiterating that manually specifying gradients is not necessary for this simple problem, but it serves as an illustrative example in the event that you do have a more complex use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.autodiffutils import AutoDiffXd, ExtractGradient, ExtractValue, InitializeAutoDiff\n",
    "\n",
    "def cost_expression(x):\n",
    "    if x.dtype == float:\n",
    "        # The cost expression with floats.\n",
    "        return x[0] + 1.5 * x[1]\n",
    "    else:\n",
    "        # Extract the values from the AutoDiffXd object\n",
    "        x_value = ExtractValue(x).flatten()\n",
    "        x_grad = ExtractGradient(x)\n",
    "\n",
    "        # Compute the cost and gradient\n",
    "        y_value = x_value[0] + 1.5 * x_value[1]\n",
    "        y_grad = np.array([1.0, 1.5]) @ x_grad\n",
    "\n",
    "        # Return the new AutoDiffXd object with the cost and gradient\n",
    "        return AutoDiffXd(y_value, y_grad)\n",
    "\n",
    "def constraint_expression(x):\n",
    "    if x.dtype == float:\n",
    "        # The constraint expression with floats\n",
    "        return np.array([\n",
    "            x[0]**2 + x[1]**2,\n",
    "            (x[0] - 0.75)**2 + (x[1] + 0.15)**2,\n",
    "            (x[0] - 0.55)**2 + (x[1] - 0.06)**2,\n",
    "        ])\n",
    "    else:\n",
    "        # Extract the values from the AutoDiffXd object\n",
    "        x_value = ExtractValue(x).flatten()\n",
    "        x_grad = ExtractGradient(x)\n",
    "\n",
    "        # Compute the constraint and gradient (with chain rule)\n",
    "        y_value = np.array([\n",
    "            x_value[0]**2 + x_value[1]**2,\n",
    "            (x_value[0] - 0.75)**2 + (x_value[1] + 0.15)**2,\n",
    "            (x_value[0] - 0.55)**2 + (x_value[1] - 0.06)**2,\n",
    "        ])\n",
    "        y_grad = 2.0 * np.array([\n",
    "            [x_value[0], x_value[1]],\n",
    "            [(x_value[0] - 0.75), (x_value[1] + 0.15)],\n",
    "            [(x_value[0] - 0.55), (x_value[1] - 0.06)],\n",
    "        ]) @ x_grad\n",
    "\n",
    "        # Return the new AutoDiffXd object with the constraint and gradient\n",
    "        return InitializeAutoDiff(y_value, y_grad)\n",
    "\n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "prog.AddCost(cost_expression, x)\n",
    "prog.AddBoundingBoxConstraint(0.0, 1.0, x)\n",
    "\n",
    "# For an expression-based constraint, the lower and upper bounds must also be numpy arrays\n",
    "lb = np.array([0.5**2, 0.25**2, 0.1**2])\n",
    "ub = np.array([np.inf, np.inf, np.inf])\n",
    "prog.AddConstraint(constraint_expression, lb, ub, x)\n",
    "\n",
    "result = Solve(prog, np.array([0.5, 0.5]))\n",
    "print(\"Success? \", result.is_success())\n",
    "print(result.get_solution_result())\n",
    "print(\"x*= \", result.GetSolution(x))\n",
    "print(\"cost = \", result.get_optimal_cost())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating Gradients using Finite Differencing\n",
    "In other cases, you may not have analytical gradients to provide to the solver.\n",
    "\n",
    "You can always rely on *finite differencing* to numerically approximate the gradients by evaluating the expression at the nominal value, as well as at variable values slightly perturbed from the nominal values.\n",
    "\n",
    "Please note that you should only rely on finite differencing when absolutely necessary, as evaluating your cost and constraint functions multiple times can drastically slow down the optimization process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.math import ComputeNumericalGradient, NumericalGradientMethod, NumericalGradientOption\n",
    "\n",
    "def cost_fn(x):\n",
    "    # This must return a numpy array to be compatible with numerical gradient extraction\n",
    "    return np.array([x[0] + 1.5 * x[1]])\n",
    "\n",
    "def cost_expression(x):\n",
    "    if x.dtype == float:\n",
    "        # The cost expression with floats.\n",
    "        return x[0] + 1.5 * x[1]\n",
    "    else:\n",
    "        # Extract the values from the AutoDiffXd object\n",
    "        x_value = ExtractValue(x)\n",
    "        x_grad = ExtractGradient(x)\n",
    "\n",
    "        # Compute the cost and approximate gradient\n",
    "        # We are using forward differencing here as the `option` argument.\n",
    "        y_value = cost_fn(x_value)\n",
    "        y_grad = ComputeNumericalGradient(\n",
    "            cost_fn, x_value, option=NumericalGradientOption(NumericalGradientMethod.kForward)\n",
    "        ) @ x_grad\n",
    "\n",
    "        # Return the new AutoDiffXd object with the cost and gradient\n",
    "        return AutoDiffXd(y_value[0,0], y_grad[0,:])\n",
    "\n",
    "def constraint_fn(x):\n",
    "    return np.array([\n",
    "        x[0]**2 + x[1]**2,\n",
    "        (x[0] - 0.75)**2 + (x[1] + 0.15)**2,\n",
    "        (x[0] - 0.55)**2 + (x[1] - 0.06)**2,\n",
    "    ])\n",
    "\n",
    "def constraint_expression(x):\n",
    "    if x.dtype == float:\n",
    "        # The constraint expression with floats\n",
    "        return constraint_fn(x)\n",
    "    else:\n",
    "        # Extract the values from the AutoDiffXd object\n",
    "        x_value = ExtractValue(x)\n",
    "        x_grad = ExtractGradient(x)\n",
    "\n",
    "        # Compute the constraint and approximate gradient\n",
    "        # We are using central differencing here as the `option` argument.\n",
    "        y_value = constraint_fn(x_value)\n",
    "        y_grad = ComputeNumericalGradient(\n",
    "            constraint_fn, x_value, option=NumericalGradientOption(NumericalGradientMethod.kCentral)\n",
    "        ) @ x_grad\n",
    "        \n",
    "        # Return the new AutoDiffXd object with the constraint and gradient\n",
    "        return InitializeAutoDiff(y_value, y_grad)\n",
    "                                  \n",
    "prog = MathematicalProgram()\n",
    "x = prog.NewContinuousVariables(2)\n",
    "prog.AddCost(cost_expression, x)\n",
    "prog.AddBoundingBoxConstraint(0.0, 1.0, x)\n",
    "\n",
    "# For an expression-based constraint, the lower and upper bounds must also be numpy arrays\n",
    "lb = np.array([0.5**2, 0.25**2, 0.1**2])\n",
    "ub = np.array([np.inf, np.inf, np.inf])\n",
    "prog.AddConstraint(constraint_expression, lb, ub, x)\n",
    "\n",
    "result = Solve(prog, np.array([0.5, 0.5]))\n",
    "print(\"Success? \", result.is_success())\n",
    "print(result.get_solution_result())\n",
    "print(\"x*= \", result.GetSolution(x))\n",
    "print(\"cost = \", result.get_optimal_cost())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
